import os
import uuid
import time
import logging
from datetime import datetime
from abc import ABC, ABCMeta, abstractmethod
from gtts import gTTS
import ChatTTS
import torch
import torchaudio

logger = logging.getLogger(__name__)


class AbstractTTS(ABC):
    __metaclass__ = ABCMeta

    @abstractmethod
    def to_tts(self, text):
        pass


class GTTS(AbstractTTS):
    def __init__(self, config):
        self.output_file = config.get("output_file")
        self.lang = config.get("lang")

    def _generate_filename(self, extension=".aiff"):
        return os.path.join(self.output_file, f"tts-{datetime.now().date()}@{uuid.uuid4().hex}{extension}")

    def _log_execution_time(self, start_time):
        end_time = time.time()
        execution_time = end_time - start_time
        logger.debug(f"执行时间: {execution_time:.2f} 秒")

    def to_tts(self, text):
        tmpfile = self._generate_filename(".aiff")
        try:
            start_time = time.time()
            tts = gTTS(text=text, lang=self.lang)
            tts.save(tmpfile)
            self._log_execution_time(start_time)
            return tmpfile
        except Exception as e:
            logger.debug(f"生成TTS文件失败: {e}")
            return None



class CHATTTS(AbstractTTS):
    def __init__(self, config):
        self.output_file = config.get("output_file", ".")
        self.spk_emb_file = config.get("spk_emb_file")  # 新增 spk_emb_file 参数
        self.lang = config.get("lang", "zh-cn")
        self.chat = ChatTTS.Chat()
        self.chat.load(compile=False)  # Set to True for better performance
        if self.spk_emb_file and os.path.exists(self.spk_emb_file):
            # 从文件加载说话人 embedding
            self.spk = torch.load(self.spk_emb_file)
        else:
            # 没有指定文件则随机采样
            self.spk = self.chat.sample_random_speaker()
            if self.spk_emb_file:
                logger.warning(f"spk_emb_file {self.spk_emb_file} 不存在，将使用随机说话人")
        # self.rand_spk = self.chat.sample_random_speaker()

    def _generate_filename(self, extension=".wav"):
        return os.path.join(self.output_file, f"tts-{datetime.now().date()}@{uuid.uuid4().hex}{extension}")

    def _log_execution_time(self, start_time):
        end_time = time.time()
        execution_time = end_time - start_time
        logger.debug(f"Execution Time: {execution_time:.2f} seconds")

    def to_tts(self, text):
        tmpfile = self._generate_filename(".wav")
        start_time = time.time()
        try:
            params_infer_code = ChatTTS.Chat.InferCodeParams(
                spk_emb=self.spk,  # add sampled speaker
                temperature=.3,  # using custom temperature
                top_P=0.7,  # top P decode
                top_K=20,  # top K decode
            )
            params_refine_text = ChatTTS.Chat.RefineTextParams(
                prompt='[oral_5]',
            )
            wavs = self.chat.infer(
                [text],
                params_refine_text=params_refine_text,
                params_infer_code=params_infer_code,
            )
            try:
                torchaudio.save(tmpfile, torch.from_numpy(wavs[0]).unsqueeze(0), 24000)
            except:
                torchaudio.save(tmpfile, torch.from_numpy(wavs[0]), 24000)
            self._log_execution_time(start_time)
            return tmpfile
        except Exception as e:
            logger.error(f"Failed to generate TTS file: {e}")
            return None


def create_instance(class_name, *args, **kwargs):
    # 获取类对象
    cls = globals().get(class_name)
    if cls:
        # 创建并返回实例
        return cls(*args, **kwargs)
    else:
        raise ValueError(f"Class {class_name} not found")
    

import logging
if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG)

    # 语音文件保存路径
    output_path = "./tts_outputs"
    os.makedirs(output_path, exist_ok=True)

    # speaker embedding 文件夹
    speaker_folder = "./tts/speaker"
    speaker_files = [f for f in os.listdir(speaker_folder) if f.endswith(".pt")]

    if not speaker_files:
        print("未找到任何 .pt 说话人文件")
        exit(1)

    # 遍历所有 speaker 文件
    for spk_file in speaker_files:
        config = {
            "output_file": output_path,
            "spk_emb_file": os.path.join(speaker_folder, spk_file),
            "lang": "zh-cn"
        }

        try:
            tts_instance = create_instance("CHATTTS", config)
            print(f"\n测试 speaker 文件: {spk_file}")
            print(f"实例类型: {type(tts_instance)}")

            text = f"测试语音文件: {spk_file}"
            result_file = tts_instance.to_tts(text)

            if result_file:
                print(f"TTS 文件生成成功: {result_file}")
            else:
                print("TTS 文件生成失败")

        except Exception as e:
            print(f"错误: {e}")